from pocketflow import Node, Flow, BatchNode

from markitdown import MarkItDown
import os

class DocumentLoaderNode(BatchNode):
    def prep(self, shared):
        # 获取用户输入的文档所在文件夹路径，如果用户输入为空，则使用默认路径
        document_path = input("请输入文档所在文件夹路径: ")
        if document_path == "":
            document_path = "RAG/docs"
        
        # 支持的文档文件扩展名
        supported_extensions = [
            # 文本文件
            '.txt', '.md', '.markdown', '.html', '.htm', '.xml', '.json', '.csv', 
            # 代码文件
            '.py', '.js', '.css', '.java', '.c', '.cpp', '.h', '.cs', '.php',
            # 文档文件
            '.pdf', '.doc', '.docx', '.ppt', '.pptx', '.xls', '.xlsx',
            # 其他常见格式
            '.rtf', '.odt', '.ods', '.odp'
        ]
        
        # 读取文件夹中的所有文件
        documents = []
        file_paths = []  # 保存文件路径用于调试
        
        for file in os.listdir(document_path):
            file_path = os.path.join(document_path, file)
            
            # 跳过目录
            if os.path.isdir(file_path):
                continue
                
            # 检查文件扩展名
            _, ext = os.path.splitext(file)
            ext = ext.lower()
            
            # 如果不是支持的文件类型，跳过
            if ext not in supported_extensions:
                print(f"跳过不支持的文件类型: {file}")
                continue
                
            # 对于文本文件，尝试直接读取
            if ext in ['.txt', '.md', '.markdown', '.html', '.htm', '.xml', '.json', '.csv', 
                      '.py', '.js', '.css', '.java', '.c', '.cpp', '.h', '.cs', '.php']:
                # 尝试多种编码格式
                encodings = ['utf-8', 'gbk', 'gb2312', 'gb18030', 'big5', 'latin1']
                content = None
                
                for encoding in encodings:
                    try:
                        with open(file_path, "r", encoding=encoding) as f:
                            content = f.read()
                        print(f"成功以 {encoding} 编码读取文件: {file}")
                        break  # 如果成功读取，跳出循环
                    except UnicodeDecodeError:
                        continue  # 如果解码错误，尝试下一种编码
                
                if content is None:
                    print(f"警告: 无法解码文件 {file}，尝试二进制读取")
                    try:
                        # 有些文本文件可能包含少量二进制数据，尝试忽略错误解码
                        with open(file_path, "rb") as f:
                            content = f.read().decode('utf-8', errors='ignore')
                        print(f"成功以二进制模式读取文件: {file}")
                    except Exception as e:
                        print(f"错误: 无法读取文件 {file}: {str(e)}")
                        continue
            else:
                # 对于二进制文档，我们不需要读取内容，MarkItDown会直接处理
                content = file_path  # 保存文件路径而不是内容
                file_paths.append(file_path)
                print(f"添加文档文件路径: {file}")
                documents.append(content)
                continue  # 跳到下一个文件
            
            # 移除可能导致问题的字符
            if content:
                content = self._clean_text(content)
            
            if content:
                documents.append(content)
                
        if not documents:
            print("警告: 未找到或无法读取任何支持的文档")
            return []
        
        # 保存文件路径到shared用于调试    
        shared["file_paths"] = file_paths
        return documents
    
    def _clean_text(self, text):
        """清理文本，移除嵌入的空字符和其他可能导致问题的字符"""
        # 移除空字符
        text = text.replace('\0', '')
        
        # 移除不可打印字符（但保留常见空白字符）
        printable = set(' \t\n\r\f\v' + ''.join(chr(i) for i in range(32, 127)))
        text = ''.join(c for c in text if c in printable)
        
        return text
    
    def exec(self, document):
        try:
            # 如果document是文件路径（针对二进制文档）
            if isinstance(document, str) and os.path.exists(document):
                print(f"使用MarkItDown转换文档文件: {document}")
                md = MarkItDown(enable_plugins=True)  # 对文档启用插件
                result = md.convert(document)
            else:
                # 直接转换文本内容
                md = MarkItDown(enable_plugins=False)
                result = md.convert(document)
            
            print("文档转换为markdown格式成功")
            return result.text_content
        except Exception as e:
            print(f"警告: 文档转换失败: {str(e)}")
            
            # 对于文本文档，返回原始文本
            if not os.path.exists(document):
                return document
                
            # 对于文件路径，尝试通过其他方式读取
            try:
                print(f"尝试使用备用方法提取文本: {document}")
                # 这里可以添加备用的文本提取方法
                # 例如使用PyPDF2、python-docx等库
                
                # 目前仅返回文件名作为标记
                filename = os.path.basename(document)
                return f"# {filename}\n\n无法转换此文档格式。"
            except Exception as e2:
                print(f"备用文本提取也失败: {str(e2)}")
                return f"# 无法读取的文档\n\n无法处理的文档格式或损坏的文件。"
    
    def post(self, shared, prep_res, exec_res_list):
        # 将markdown格式内容存储到shared中
        all_markdown_content = []
        for exec_res in exec_res_list:
            if exec_res:  # 只添加非空内容
                all_markdown_content.append(exec_res)
        
        if not all_markdown_content:
            print("错误: 所有文档处理后均为空")
            all_markdown_content = ["# 示例文档\n这是一个示例文档，因为未能成功加载任何真实文档。"]
            
        shared["all_markdown_content"] = all_markdown_content
        print("markdown格式内容存储到shared中all_markdown_content成功")
        return "default"
        

from utils import chunk_markdown

class ChunkDocumentNode(BatchNode):
    def prep(self, shared):
        documents = shared["all_markdown_content"]
        return documents
    
    def exec(self, document):
        # 将markdown格式的内容进行切片
        chunks = chunk_markdown(document)
        return chunks
    
    def post(self, shared, prep_res, exec_res_list):
        # 将切片后的内容存储到shared中
        all_chunks = []
        for exec_res in exec_res_list:
            all_chunks.extend(exec_res)
        shared["all_chunks"] = all_chunks
        print("切片后的内容存储到shared中all_chunks成功")
        return "default"
    
from utils import get_zhipuai_embedding
class EmbeddingNode(BatchNode):
    def prep(self, shared):
        chunks = shared["all_chunks"]
        return chunks
    
    def exec(self, chunk):
        # 使用embedding模型对内容进行向量化
        try:
            embeddings = get_zhipuai_embedding(chunk)
            return embeddings
        except Exception as e:
            print(f"向量化失败: {str(e)}")
            # 返回零向量作为后备方案
            return None
    
    def post(self, shared, prep_res, exec_res_list):
        # 将向量化后的内容存储到shared中
        all_embeddings = []
        for exec_res in exec_res_list:
            if exec_res is not None:  # 只添加成功的嵌入
                all_embeddings.append(exec_res)
        
        if not all_embeddings:
            print("错误: 所有文本向量化后均为空")
            # 创建一个示例向量，维度为1536（智谱AI embedding的默认维度）
            shared["all_embeddings"] = np.zeros((1, 1536), dtype=np.float32)
        else:
            # 将列表转换为NumPy数组
            shared["all_embeddings"] = np.array(all_embeddings, dtype=np.float32)
            
        print(f"向量化后的内容存储到shared中all_embeddings成功，形状: {shared['all_embeddings'].shape}")
        return "default"
    
import faiss

class CreateIndexNode(Node):
    def prep(self, shared):
        """获取向量化后的内容和文本"""
        all_embeddings = shared["all_embeddings"]
        texts = shared.get("all_chunks", [])
        return all_embeddings, texts
    
    def exec(self, inputs):
        """创建FAISS索引并添加向量"""
        embeddings, texts = inputs
        print("🔍 Creating search index...")
        print(f"嵌入向量形状: {embeddings.shape}")
        
        dimension = embeddings.shape[1]
        
        # 创建一个L2范数索引
        index = faiss.IndexFlatL2(dimension)
        
        # 添加向量到索引
        index.add(embeddings)
        
        # 检查文本和向量的数量是否匹配
        if len(texts) != embeddings.shape[0]:
            print(f"警告: 文本数量({len(texts)})与向量数量({embeddings.shape[0]})不匹配")
            # 如果不匹配，进行调整
            if len(texts) < embeddings.shape[0]:
                # 用占位符填充文本
                texts = texts + ["[占位文本]"] * (embeddings.shape[0] - len(texts))
            else:
                # 截断文本
                texts = texts[:embeddings.shape[0]]
        
        return index, texts
    
    def post(self, shared, prep_res, exec_res):
        """将索引和文本存储到shared中"""
        index, texts = exec_res
        shared["index"] = index
        shared["texts"] = texts
        print(f"✅ Index created with {index.ntotal} vectors")
        return "default"

import numpy as np

# Nodes for the online flow
class EmbedQueryNode(Node):
    def prep(self, shared):
        """获取用户输入的查询内容"""
        query = shared.get("query", "")
        if query == "":
            query = input("请输入查询内容: ")
            # 将用户输入的查询内容保存到shared中
            shared["query"] = query
            
        if not query.strip():
            raise ValueError("查询内容为空")
        return query
    
    def exec(self, query):
        """使用embedding模型对查询内容进行向量化"""
        print(f"🔍 向量化查询: {query}")
        try:
            query_embedding = get_zhipuai_embedding(query)
            # 确保返回的是1维向量，然后转换为2维向量用于FAISS
            return np.array([query_embedding], dtype=np.float32)
        except Exception as e:
            print(f"查询向量化失败: {str(e)}")
            # 创建一个零向量作为后备方案，维度为1536（智谱AI embedding的默认维度）
            return np.zeros((1, 1536), dtype=np.float32)
    
    def post(self, shared, prep_res, exec_res):
        """将向量化后的查询内容存储到shared中"""
        shared["query_embedding"] = exec_res
        print(f"查询向量化完成，形状: {exec_res.shape}")
        return "default"
    

class RetrieveDocumentNode(Node):
    def prep(self, shared):
        """获取向量化后的查询内容、FAISS索引和文本内容"""
        query_embedding = shared.get("query_embedding")
        index = shared.get("index")
        texts = shared.get("texts", [])
        
        # 验证所有必要的组件都存在
        if query_embedding is None:
            raise ValueError("查询向量未找到")
        if index is None:
            raise ValueError("索引未找到")
        if not texts:
            raise ValueError("文本内容为空")
            
        return query_embedding, index, texts
    
    def exec(self, inputs):
        """在索引中搜索最相似的文档"""
        print("🔎 搜索最相似的文档...")
        query_embedding, index, texts = inputs
        
        # 安全检查
        if index.ntotal == 0:
            print("警告: 索引为空")
            return {
                "text": "无相关文档",
                "index": -1,
                "distance": float('inf')
            }
            
        # 检查文本和索引数量是否匹配
        if len(texts) < index.ntotal:
            print(f"警告: 文本数量({len(texts)})小于索引向量数量({index.ntotal})")
            # 创建一个新的文本列表，而不是修改原始列表
            texts = texts + ["[缺失文本]"] * (index.ntotal - len(texts))
        
        # 在索引中搜索最相似的文档
        k = min(5, index.ntotal)  # 搜索前k个结果，但不超过索引总数
        distances, indices = index.search(query_embedding, k=k)
        
        # 获取最相似文档的索引
        results = []
        for i in range(k):
            idx = indices[0][i]
            distance = distances[0][i]
            
            if idx >= 0 and idx < len(texts):
                text = texts[idx]
                results.append({
                    "text": text,
                    "index": int(idx),
                    "distance": float(distance)
                })
        
        # 如果没有找到任何结果，返回默认值
        if not results:
            return {
                "text": "无相关文档",
                "index": -1,
                "distance": float('inf')
            }
            
        # 返回最相似的文档
        return results[0]
    
    def post(self, shared, prep_res, exec_res):
        """将检索到的文档存储到shared中"""
        shared["retrieved_document"] = exec_res
        print(f"📄 检索到的文档 (index: {exec_res['index']}, distance: {exec_res['distance']:.4f})")
        print(f"📄 最相似的文本: \"{exec_res['text'][:100]}...\"")  # 只显示前100个字符
        return "default"
    
from utils import call_zhipuai
class GenerateAnswerNode(Node):
    def prep(self, shared):
        """获取用户输入的查询内容、检索到的文档和任何其他需要的上下文"""
        query = shared.get("query", "")
        retrieved_doc = shared.get("retrieved_document", {"text": "无相关文档", "index": -1, "distance": float('inf')})
        
        # 检查查询是否为空
        if not query.strip():
            raise ValueError("查询内容为空")
            
        return query, retrieved_doc
    
    def exec(self, inputs):
        """使用LLM生成答案"""
        query, retrieved_doc = inputs
        
        # 检查检索到的文档是否有效
        doc_text = retrieved_doc.get("text", "")
        distance = retrieved_doc.get("distance", float('inf'))
        
        # 如果距离太大，可能不相关
        distance_warning = ""
        if distance > 100:  # 根据您的数据调整这个阈值
            distance_warning = "注意：检索到的文档可能与问题不太相关。"
        
        prompt = f"""
根据提供的上下文回答以下问题：

问题: {query}

上下文: {doc_text}

{distance_warning}

请直接回答问题，不要复述问题或上下文。如果上下文中没有相关信息，请明确说明"根据提供的上下文无法回答这个问题"。

答案:
"""
        
        try:
            answer = call_zhipuai([
                {"role": "user", "content": prompt}
            ])
            return answer
        except Exception as e:
            print(f"LLM调用失败: {str(e)}")
            return f"抱歉，生成回答时出现错误: {str(e)}"
    
    def post(self, shared, prep_res, exec_res):
        """将生成的答案存储到shared中"""
        shared["generated_answer"] = exec_res
        print("\n🤖 生成的答案:")
        print("-" * 50)
        print(exec_res)
        print("-" * 50)
        return "default"